package com.apop.demo.encrypt.util.sm;


import com.apop.demo.encrypt.util.Base64Utils;
import com.apop.demo.encrypt.util.HexUtils;
import com.apop.demo.encrypt.util.sm.cipher.SMKeyPair;

import org.bouncycastle.asn1.gm.GMNamedCurves;
import org.bouncycastle.asn1.x9.ECNamedCurveTable;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.KeyGenerationParameters;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.generators.ECKeyPairGenerator;
import org.bouncycastle.crypto.params.*;

import org.bouncycastle.math.ec.ECPoint;


import java.math.BigInteger;
import java.security.SecureRandom;


public class FzSM2 {
    private static final String ECNAME = "sm2p256v1";

    private static final String RANDOMSOURCE = "SHA1PRNG";

    public static byte[] encrypt(byte[] publicKey, byte[] data) throws InvalidCipherTextException {
        X9ECParameters sm2ECParameters = ECNamedCurveTable.getByName("sm2p256v1");
        ECDomainParameters domainParameters = new ECDomainParameters(sm2ECParameters.getCurve(), sm2ECParameters.getG(),
                sm2ECParameters.getN());
        SM2Engine sm2Engine = new SM2Engine();
        ECPoint pukPoint = sm2ECParameters.getCurve().decodePoint(publicKey);
        ECPublicKeyParameters publicKeyParameters = new ECPublicKeyParameters(pukPoint, domainParameters);
        ParametersWithRandom paramRandom = new ParametersWithRandom((CipherParameters) publicKeyParameters);
        sm2Engine.init(true, (CipherParameters) paramRandom);
        return sm2Engine.processBlock(data, 0, data.length);
    }

    public static byte[] encrypt(String publicKey, int keyType, byte[] data) throws InvalidCipherTextException {
        byte[] publicKeyBytes = null;
        if (keyType == 1) {
            publicKeyBytes = HexUtils.hexToByte(publicKey);
        } else {
            if (keyType != 2) {
                throw new IllegalArgumentException("keyType parameter is illegal! ");
            }
            publicKeyBytes = Base64Utils.bouncyDecode(publicKey);
        }
        return encrypt(publicKeyBytes, data);
    }

    public static byte[] decrypt(byte[] privateKey, byte[] data) throws InvalidCipherTextException {
        X9ECParameters sm2ECParameters = ECNamedCurveTable.getByName("sm2p256v1");
        ECDomainParameters domainParameters = new ECDomainParameters(sm2ECParameters.getCurve(), sm2ECParameters.getG(),
                sm2ECParameters.getN());
        SM2Engine sm2Engine = new SM2Engine();
        BigInteger privateKeyD = new BigInteger(1, privateKey);
        ECPrivateKeyParameters privateKeyParameters = new ECPrivateKeyParameters(privateKeyD, domainParameters);
        sm2Engine.init(false, (CipherParameters) privateKeyParameters);
        return sm2Engine.processBlock(data, 0, data.length);
    }

    public static byte[] decrypt(String privateKey, int keyType, byte[] data) throws InvalidCipherTextException {
        byte[] privateKeyBytes = null;
        if (keyType == 1) {
            privateKeyBytes = HexUtils.hexToByte(privateKey);
        } else {
            if (keyType != 2) {
                throw new IllegalArgumentException("keyType parameter is illegal! ");
            }
            privateKeyBytes = Base64Utils.bouncyDecode(privateKey);
        }
        return decrypt(privateKeyBytes, data);
    }

    private static AsymmetricCipherKeyPair genKeyPair0() throws Exception {
        X9ECParameters sm2ECParameters = GMNamedCurves.getByName("sm2p256v1");
        ECDomainParameters domainParameters = new ECDomainParameters(sm2ECParameters.getCurve(), sm2ECParameters.getG(),
                sm2ECParameters.getN());
        ECKeyPairGenerator keyPairGenerator = new ECKeyPairGenerator();
        keyPairGenerator.init((KeyGenerationParameters) new ECKeyGenerationParameters(domainParameters,
                SecureRandom.getInstance("SHA1PRNG")));
        return keyPairGenerator.generateKeyPair();
    }

    public static SMKeyPair genKeyPair(int keyType) throws Exception {
        AsymmetricCipherKeyPair asymmetricCipherKeyPair = genKeyPair0();
        ECPoint ecPoint = ((ECPublicKeyParameters) asymmetricCipherKeyPair.getPublic()).getQ();
        byte[] pubKey = ecPoint.getEncoded(false);
        BigInteger priKey = ((ECPrivateKeyParameters) asymmetricCipherKeyPair.getPrivate()).getD();
        String pubKeyStr = null;
        String priKeyStr = null;
        if (keyType == 1) {
            pubKeyStr = HexUtils.byteToHex(pubKey).toUpperCase();
            priKeyStr = priKey.toString(16).toUpperCase();
        } else {
            if (keyType != 2) {
                throw new IllegalArgumentException("keyType parameter is illegal! ");
            }
            pubKeyStr = new String(Base64Utils.bouncyEncode(pubKey));
            priKeyStr = new String(Base64Utils.bouncyEncode(priKey.toByteArray()));
        }
        return new SMKeyPair(priKeyStr, pubKeyStr);
    }

    public static void main(String[] args) throws Exception {
        SMKeyPair k = genKeyPair(1);
        System.out.println(k.getPriKey());
        System.out.println(k.getPubKey());
    }
}
